import csv

def logging(file):
    def write_log(s):
        print(s)
        with open(file, 'a') as f:
            f.write(s)

    return write_log


def logging_csv(file):
    def write_csv(s):
        with open(file, 'a', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(s)

    return write_csv

def calc_mean_avg_prec(preds):
    mean_avg_prec, relQ = 0., 0.
    for pred in preds.values():
        sorted_pred = sorted(pred, key=lambda res: res[1])
        sorted_pred = sorted(sorted_pred, key=lambda res: res[2], reverse=True)
        avg_prec, rel = 0., 0.
        for i,tri in enumerate(sorted_pred):
            if tri[1] == 1:
                rel += 1.
                avg_prec += rel/(i+1)
        avg_prec /= rel
        mean_avg_prec += avg_prec
        relQ += 1.
    mean_avg_prec /= relQ
    return mean_avg_prec

def MAP(labels, preds):
    """
    labels: binary labels
    preds: predicted scores, the order of labels and preds are the same
    """
    total_score = 0.
    for label, pred in zip(labels, preds):
        # pred is used for sorting
        scores = [(i, pred[i], label[i]) for i in range(len(pred))]
        scores.sort(key=lambda k:k[1], reverse=True)

        score = 0.
        num_hits = 0.
        for i in range(len(scores)):
            if scores[i][2] == 1:
                num_hits += 1.
                score += num_hits/(i+1.)
        total_score += score / max(1., sum(label))

    return total_score/float(len(labels))

def average_precision(gt, pred):
    if not gt:
        return 0.0

    score = 0.
    num_hits = 0.
    for i,p in enumerate(pred):
        if p in gt and p not in pred[:i]:
            num_hits += 1.
            score += num_hits/(i+1.)

    return score / max(1., len(gt))

def RR(gt, pred):
    score = 0,
    for rank, item in enumerate(pred):
        if item in gt:
            score = 1./(rank+1.)
            break
    return score

def MRR(labels, preds):
    total_score = 0.
    for label, pred in zip(labels, preds):
        # pred is used for sorting
        scores = [(i, pred[i], label[i]) for i in range(len(pred))]
        scores.sort(key=lambda k:k[1], reverse=True)

        score = 0.
        for i in range(len(scores)):
            if scores[i][2] == 1:
                score = 1./(i+1.)
                break
        total_score += score

    return total_score/float(len(labels))
